package chapter30;

import java.net.*;
import java.io.*;
import java.util.*;

public class DnsClient {
	private static final int DNS_PORT = 53;
	private static final int REQ_MAX_LEN  = 512;
	private static final int RESP_MAX_LEN = 512;
	private static final String[] RCODES = {"Success (0)", "Format error (1)", "Server error (2)", 
											"Name error (3)", "Not implemented (4)", "Denied (5)"};
	private static final int A = 1, NS = 2, CNAME = 5, PTR = 12, MX = 15, AAAA = 28, ANY = 256;

	//
	// Main
	//
	public static void main(String[] args) {
		if (args.length < 2) {
			System.err.println("usage: java DnsClient server query");
			System.exit(0);
		}
		
		DnsClient dnsClient = new DnsClient();
		dnsClient.query(args[0], args[1]);
	}
	
	public void query(String server, String query) {
		DatagramSocket socket = null;
		Random rand = new Random();

		try {
			socket = new DatagramSocket();
			int id = (short)rand.nextInt();
			
			request(socket, id, server, query);
			response(socket, id);

		} catch (SocketException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} finally {
			socket.close();
		}
	}
	
	//
	// DNS query request
	//
	private void request(DatagramSocket socket, int id, String server, String query) throws IOException {
		byte[] buf = new byte[REQ_MAX_LEN];
		
		// (1) id
		buf[0] = (byte)((id >> 8) & 0xff);
		buf[1] = (byte)(id & 0xff);
		// (2) flags
		buf[2] = 0x01; // query, recursion desired
		buf[3] = 0x00;
		// (3) # of queries
		buf[4] = 0x00; //
		buf[5] = 0x01; // a single query
		// (4) # of response RRs
		// (5) # of authoritative RRs
		// (6) # of additional RRs
		Arrays.fill(buf, 6, 12, (byte)0);
		
		// (7) query records
        int qlen = query.length();
        String[] qarr = query.split("\\.");
        int rlen = 0;
        
        for (int i = 0; i < qarr.length; i++)
        {
            buf[12 + rlen] = (byte) qarr[i].length();
            byte[] subq = qarr[i].getBytes();
            for (int j = 1; j <= subq.length; j++)
            {
                buf[12 + rlen + j] = subq[j - 1];
            }
            
            rlen += subq.length + 1; 
        }
        
        buf[13 + qlen] = 0x00;
		//
		//
		// DONE: make query string
		// (e.g.) www.yahoo.co.jp -> 3www5yahoo2co2jp\0
		//
		//
		buf[14 + qlen] = 0x00; // query type: A
		buf[15 + qlen] = 0x01; // 
		buf[16 + qlen] = 0x00; // the Internet (IN)
		buf[17 + qlen] = 0x01; // 
		
		InetAddress addr = InetAddress.getByName(server);
		//
		//
		// DONE: send a packet
		//
		//
        DatagramPacket sendPacket = new DatagramPacket(buf, buf.length, addr, DNS_PORT);
        socket.send(sendPacket);
	}
	
	//
	// DNS response
	//
	private void response(DatagramSocket socket, int reqId) throws IOException {
		byte[] buf = new byte[RESP_MAX_LEN];
		int length;

		//
		//
		// DONE: receive a packet
		//
		//
        DatagramPacket receivePacket = new DatagramPacket(buf, buf.length);
        socket.receive(receivePacket);
		length = receivePacket.getLength();
        
		parseResponse(buf, length, reqId);
	}
	
	private void parseResponse(byte[] buf, int length, int reqId) {
		System.out.println("RESPONSE:");
		
		// (1) id
		int id = (buf[0] & 0xff) << 8 | (buf[1] & 0xff);
		System.out.println("Response ID=" + (short)id + ", Request ID=" + (short)reqId);
		// (2) flags
		int flags = (buf[2] << 8) + buf[3];
		System.out.println("Query/Response=" + ((flags >> 15) & 0x01));
		System.out.println("OPCODE=" + ((flags >> 11) & 0x0f));
		System.out.println("Authoritative Answer=" + ((flags >> 10) & 0x01));
		System.out.println("Truncation=" + ((flags >> 9) & 0x01));
		System.out.println("Recursion Desired=" + ((flags >> 8) & 0x01));
		System.out.println("Recursion Available=" + ((flags >> 7) & 0x01));
		System.out.println("Response Code=" + RCODES[(flags & 0x0f)]);
		// (3) # of queries
		int nq = (buf[4] << 8) + buf[5];
		System.out.println("Number of queries=" + nq);
		// (4) # of response RRs
		int respRRs = (buf[6] << 8) + buf[7];
		System.out.println("Number of response RRs=" + respRRs);
		// (5) # of authoritative RRs
		int authRRs = (buf[8] << 8) + buf[9];
		System.out.println("Number of authoritative RRs=" + authRRs);
		// (6) # of additional RRs
		int addRRs = (buf[10] << 8) + buf[11];
		System.out.println("Number of additional RRs=" + addRRs);
		System.out.println();
		
		// (7) query records
		System.out.println("QUESTION SECTION:");
		int offset = 12;
		offset = queryRecord(buf, offset);
		offset = queryTypeClass(buf, offset, false);
		System.out.println();
		
		// (8-10) resource records
		System.out.println("ANSWER SECTION:");
		while ((offset = readRecord(buf, offset)) != length) System.out.println();
	}
	
	private int readRecord(byte[] buf, int offset) {
		// (8) response records
		offset = queryRecord(buf, offset);
		offset = queryTypeClass(buf, offset, true);
		return offset;
	}
	
	private int queryRecord(byte[] buf, int offset) {
		while (buf[offset] != 0x00) {
			if (((buf[offset] >> 6) & 0x03) == 0x03) {
				int ptr = ((buf[offset++] << 8) & 0x3f) + (buf[offset++] & 0xff);
                
				queryRecord(buf, ptr);
				return offset;
				
			} else {
				int num = buf[offset++];
				for (int i = 0; i < num; i++) { System.out.print((char)buf[offset++]); }
				System.out.print('.');
			}
		}
		System.out.println();
		return ++offset;
	}
	
	private int queryTypeClass(byte[] buf, int offset, boolean resRecord) {
		int qtype = (buf[offset++] << 8) + buf[offset++];
		System.out.print("Type=");
		switch (qtype) {
			case A:     System.out.println("A"); break;
			case NS:    System.out.println("NS"); break;
			case CNAME: System.out.println("CNAME"); break;
			case PTR:   System.out.println("PTR"); break;
			case MX:    System.out.println("MX"); break;
			case AAAA:  System.out.println("AAAA"); break;
			case ANY:   System.out.println("ANY"); break;
			default:    System.out.println("Unknown");
		}
		int cls = (buf[offset++] << 8) + buf[offset++];
		System.out.println("Class=" + ((cls == 0x01) ? "IN" : "Unknown"));
		
		if (resRecord) {
			int ttl = (buf[offset++] & 0xff) << 24 | (buf[offset++] & 0xff) << 16 | (buf[offset++] & 0xff) << 8 | (buf[offset++] & 0xff);
			System.out.println("TTL=" + ttl);
		
			int len = (buf[offset++] << 8) + buf[offset++];
			System.out.println("Resource length=" + len);
			switch (qtype) {
				case A:
					for (int i = 0; i < len; i++) {
						System.out.print(buf[offset + i] & 0xff);
						if (i != len - 1) System.out.print(".");
					}
					System.out.println(); 
					break;
					
				case NS:
					queryRecord(buf, offset);
					break;
				case CNAME:   
					queryRecord(buf, offset);
					break;
				case PTR:  
					queryRecord(buf, offset);
					break;
					
				case MX:  
					int priority = (buf[offset] << 8) + buf[offset + 1];
					System.out.println("Priority=" + priority);
					queryRecord(buf, offset + 2);
					break;
					
				case AAAA:  
					for (int i = 0; i < len; i++) {
						System.out.print(buf[offset + i] & 0xff);
						if (i != len - 1) System.out.print(":");
					}
					System.out.println(); 
					break;
					
				case ANY: 
					System.out.println("ANY"); 
					break;
				default:  
					System.out.println("Unknown");
			}
			offset += len;
		}
		
		return offset;
	}
}
